import pandas as pd
import numpy as np

class MyGBDT:
    def __init__(self, n_trees=100, learning_rate=0.1, max_depth=3):
        self.n_trees = n_trees
        self.learning_rate = learning_rate
        self.max_depth = max_depth
        self.trees = []

    def fit(self, X, y):
        # 初始化预测值
        current_pred = np.full(y.shape, np.mean(y))

        for _ in range(self.n_trees):
            # 计算负梯度（伪残差）
            residuals = y - current_pred

            # 训练新树拟合残差
            tree = DecisionTree(max_depth=self.max_depth)
            tree.fit(X, residuals)

            # 更新预测
            current_pred += self.learning_rate * tree.predict(X)

            self.trees.append(tree)

class TreeNode:
    def __init__(self, depth):
        self.depth = depth
        self.split_feature = None
        self.split_value = None
        self.left = None
        self.right = None
        self.leaf_value = None

def find_best_split(X, gradients, hessians):
    best_gain = -float('inf')
    best_feature, best_value = None, None

    # 遍历所有特征
    for feature in range(X.shape[1]):
        # 直方图构建（核心优化点）
        histogram = build_histogram(X[:, feature], gradients, hessians)

        # 寻找最佳分裂点
        for bin in range(histogram.n_bins):
            gain = calculate_split_gain(histogram, bin)
            if gain > best_gain:
                best_gain = gain
                best_feature = feature
                best_value = histogram.bin_edges[bin]

    return best_feature, best_value


class MyLightGBM:
    def __init__(self, n_bins=256):
        self.n_bins = n_bins
        self.bin_edges = None
        self.grad_sum = np.zeros(n_bins)
        self.hess_sum = np.zeros(n_bins)


def build_histogram(feature_values, gradients, hessians):
    hist = MyLightGBM()

    # 等宽分箱
    hist.bin_edges = np.linspace(feature_values.min(), feature_values.max(), hist.n_bins)

    # 将样本分配到bin
    bin_indices = np.digitize(feature_values, hist.bin_edges)

    # 累加梯度统计量
    for i in range(len(feature_values)):
        bin_idx = bin_indices[i]
        hist.grad_sum[bin_idx] += gradients[i]
        hist.hess_sum[bin_idx] += hessians[i]

    return hist


def grow_tree(node, X, gradients, hessians, depth):
    if depth >= max_depth:
        node.leaf_value = calculate_leaf_value(gradients, hessians)
        return

    # 寻找最佳分裂
    feature, value = find_best_split(X, gradients, hessians)

    if feature is None:  # 无法继续分裂
        node.leaf_value = calculate_leaf_value(gradients, hessians)
        return

    # 分裂节点
    left_idx = X[:, feature] <= value
    right_idx = X[:, feature] > value

    node.split_feature = feature
    node.split_value = value

    # 仅扩展增益最大的叶子（Leaf-wise）
    node.left = TreeNode(depth + 1)
    node.right = TreeNode(depth + 1)

    grow_tree(node.left, X[left_idx], gradients[left_idx], hessians[left_idx], depth + 1)
    grow_tree(node.right, X[right_idx], gradients[right_idx], hessians[right_idx], depth + 1)